from __future__ import print_function
import torch
import torch.nn as nn
import torch.utils.data
import torch.nn.functional as F

# 基本卷积块，输入通道，输出通道，转置卷积，是否3d卷积，批归一化，LeakyReLU，其它参数
class BasicConv(nn.Module):
    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, bn=True, relu=True, **kwargs):
        super(BasicConv, self).__init__()

        self.relu = relu
        self.use_bn = bn
        if is_3d: # 3d
            if deconv: # 转置卷积
                self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) # 3d转置卷积
            else: # 3d卷积
                self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs) # 3d卷积
            self.bn = nn.BatchNorm3d(out_channels) # 3d批归一化
        else: # 2d
            if deconv: # 转置卷积
                self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) # 2d转置卷积
            else: # 2d卷积
                self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) # 2d卷积
            self.bn = nn.BatchNorm2d(out_channels) # 2d批归一化

    def forward(self, x):
        x = self.conv(x) # 2d或3d或转置卷积
        if self.use_bn:
            x = self.bn(x) # 2d或3d批归一化
        if self.relu:
            x = nn.LeakyReLU()(x) #, inplace=True) # LeakyReLU
        return x

# 拼接块，输入通道，输出通道，转置卷积，是否3d卷积，拼接或残差连接，保持拼接，批归一化，LeakyReLU，保持视差尺寸
class Conv2x(nn.Module):
    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, bn=True, relu=True, keep_dispc=False):
        super(Conv2x, self).__init__()
        self.concat = concat
        self.is_3d = is_3d
        if deconv and is_3d: # 3d转置卷积
            kernel = (4, 4, 4) # 卷积核
        elif deconv: # 2d转置卷积
            kernel = 4 # 卷积核
        else: # 2d或3d卷积
            kernel = 3 # 卷积核

        if deconv and is_3d and keep_dispc: # 3d转置卷积，保持视差尺寸，(B,C,D,H,W)
            kernel = (1, 4, 4) # 卷积核
            stride = (1, 2, 2) # 步长
            padding = (0, 1, 1)
            self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, stride=stride, padding=padding) # 基本卷积块，3d转置卷积，尺寸*2(除视差维度)
        else: # 默认，调用：2d转置卷积，(B,C,H,W)
            self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, stride=2, padding=1) # 基本卷积块，尺寸*2，调用：2d转置卷积

        if self.concat: # 拼接，默认，调用：2d卷积，尺寸不变
            mul = 2 if keep_concat else 1 # 若保持拼接，输出通道*2，否则*1
            self.conv2 = BasicConv(out_channels*2, out_channels*mul, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1) # 基本卷积块，调用：2d卷积，尺寸不变
        else: # 残差连接
            self.conv2 = BasicConv(out_channels, out_channels, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1) # 基本卷积块，尺寸不变

    def forward(self, x, rem):
        x = self.conv1(x) # 2d转置卷积，尺寸*2
        if x.shape != rem.shape: # 两个张量尺寸不一致
            x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='nearest') # 最近邻插值x到rem的形状
        if self.concat: # 拼接，默认
            x = torch.cat((x, rem), 1) # 在通道维度连接x和rem
        else: # 残差连接
            x = x + rem
        x = self.conv2(x) # 默认拼接，调用：2d卷积，尺寸不变
        return x

def disparity_regression(x, maxdisp):
    assert len(x.shape) == 4
    disp_values = torch.arange(0, maxdisp, dtype=x.dtype, device=x.device)
    disp_values = disp_values.view(1, maxdisp, 1, 1)
    return torch.sum(x * disp_values, 1, keepdim=False)

def groupwise_difference(fea1, fea2, num_groups):
    B, G, C, H, W = fea1.shape
    cost = torch.pow((fea1 - fea2), 2).sum(2)
    assert cost.shape == (B, num_groups, H, W)
    return cost

def build_struct_volume(refimg_fea, targetimg_fea, maxdisp, num_groups):
    B, G, C, H, W = refimg_fea.shape
    volume = refimg_fea.new_zeros([B, num_groups, maxdisp, H, W])
    for i in range(maxdisp):
        if i > 0:
            volume[:, :, i, :, i:] = groupwise_difference(refimg_fea[:, :, :, :, i:], targetimg_fea[:, :, :, :, :-i], num_groups)
        else:
            volume[:, :, i, :, :] = groupwise_difference(refimg_fea, targetimg_fea, num_groups)
    volume = volume.contiguous()
    return volume

# 构建串联体积
def build_concat_volume(refimg_fea, targetimg_fea, maxdisp):
    B, C, H, W = refimg_fea.shape # 1/4原尺寸左视图上下文特征(48通道)
    volume = refimg_fea.new_zeros([B, 2 * C, maxdisp, H, W]) # 创建形状为(B,2*C,D/4,H/4,W/4)的零张量
    for i in range(maxdisp): # 对于每一个视差级别，分别串联两个张量
        if i > 0: # 左视图不变，将右视图右移i像素，串联张量
            volume[:, :C, i, :, :] = refimg_fea[:, :, :, :] # 对W维度切片，串联左视图和右视图-i开始的切片张量
            volume[:, C:, i, :, i:] = targetimg_fea[:, :, :, :-i]
        else: # 视差级别0
            volume[:, :C, i, :, :] = refimg_fea # 直接串联两个张量
            volume[:, C:, i, :, :] = targetimg_fea
    volume = volume.contiguous() # 形成串联体积，输出(B,2*C,D/4,H/4,W/4)
    return volume

# 计算组相关性
def groupwise_correlation(fea1, fea2, num_groups):
    B, C, H, W = fea1.shape # 1/4原尺寸切片左视图上下文特征(48通道)
    assert C % num_groups == 0
    channels_per_group = C // num_groups # 每个组的通道数量
    cost = (fea1 * fea2).view([B, num_groups, channels_per_group, H, W]).mean(dim=2) # 计算相关性，并在通道维度度求均值
    assert cost.shape == (B, num_groups, H, W) # 输出形状(B,G,H,W)
    return cost

# 构建组相关性体积
def build_gwc_volume(refimg_fea, targetimg_fea, maxdisp, num_groups):
    B, C, H, W = refimg_fea.shape # 1/4原尺寸左视图上下文特征(48通道)
    volume = refimg_fea.new_zeros([B, num_groups, maxdisp, H, W]) # 创建形状为(B,G,D/4,H/4,W/4)的零张量
    for i in range(maxdisp): # 对于每一个视差级别，分别计算两个张量的组相关性
        if i > 0: # 将左视图切片i像素，右视图右移i像素，计算组相关性
            volume[:, :, i, :, i:] = groupwise_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i], num_groups) # 对W维度切片，对左视图i开始和右视图-i开始的切片张量计算组相关性
        else: # 视差级别0
            volume[:, :, i, :, :] = groupwise_correlation(refimg_fea, targetimg_fea, num_groups) # 直接计算两个张量的组相关性
    volume = volume.contiguous() # 形成组相关性体积，输出(B,G,D/4,H/4,W/4)
    return volume

# 计算归一化组相关性，归一化余弦相似度
def groupwise_correlation_norm(fea1, fea2, num_groups):
    B, C, H, W = fea1.shape # 1/4原尺寸切片左视图上下文特征(48通道)
    assert C % num_groups == 0
    channels_per_group = C // num_groups # 每个组的通道数量
    fea1 = fea1.view([B, num_groups, channels_per_group, H, W]) # 分组，改变形状为(B,G,C/G,H,W)
    fea2 = fea2.view([B, num_groups, channels_per_group, H, W])
    # 在通道维度计算均值(归一化左特征*归一化右特征)，输出(B,G,H/4,W/4)；归一化操作：特征/(在通道维度计算特征的L2范数(平方和开根)+极小正值)，*表示对应元素相乘(哈达玛积)
    cost = ((fea1/(torch.norm(fea1, 2, 2, True)+1e-05)) * (fea2/(torch.norm(fea2, 2, 2, True)+1e-05))).mean(dim=2)
    assert cost.shape == (B, num_groups, H, W) # 输出形状(B,G,H,W)
    return cost

# 构建归一化组相关性体积
def build_gwc_volume_norm(refimg_fea, targetimg_fea, maxdisp, num_groups):
    B, C, H, W = refimg_fea.shape # 1/4原尺寸左视图上下文特征(48通道)
    volume = refimg_fea.new_zeros([B, num_groups, maxdisp, H, W]) # 创建形状为(B,G,D/4,H/4,W/4)的零张量
    for i in range(maxdisp): # 对于每一个视差级别，分别计算两个张量的归一化组相关性
        if i > 0: # 将左视图切片i像素，右视图右移i像素，计算归一化组相关性
            volume[:, :, i, :, i:] = groupwise_correlation_norm(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i], num_groups) # 对W维度切片，对左视图i开始和右视图-i开始的切片张量计算归一化组相关性
        else: # 视差级别0
            volume[:, :, i, :, :] = groupwise_correlation_norm(refimg_fea, targetimg_fea, num_groups) # 直接计算两个张量的归一化组相关性
    volume = volume.contiguous() # 形成归一化组相关性体积，输出(B,G,D/4,H/4,W/4)
    return volume

# 计算归一化相关性，归一化余弦相似度
def norm_correlation(fea1, fea2):
    # 在通道维度计算均值(归一化左特征*归一化右特征)，输出(B,1,H/4,W/4)；归一化操作：特征/(在通道维度计算特征的L2范数(平方和开根)+极小正值)，*表示对应元素相乘(哈达玛积)
    cost = torch.mean(((fea1/(torch.norm(fea1, 2, 1, True)+1e-05)) * (fea2/(torch.norm(fea2, 2, 1, True)+1e-05))), dim=1, keepdim=True)
    return cost

# 构建归一化相关性体积
def build_norm_correlation_volume(refimg_fea, targetimg_fea, maxdisp): # 输入1/4原尺寸左右上下文特征(48通道)，最大视差/4
    B, C, H, W = refimg_fea.shape # 1/4原尺寸左视图上下文特征(48通道)
    volume = refimg_fea.new_zeros([B, 1, maxdisp, H, W]) # 创建形状为(B,1,D/4,H/4,W/4)的零张量
    for i in range(maxdisp): # 对于每一个视差级别，分别计算两个张量的归一化相关性
        if i > 0: # 将左视图切片i像素，右视图右移i像素，计算归一化相关性
            volume[:, :, i, :, i:] = norm_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i]) # 对W维度切片，对左视图i开始和右视图-i开始的切片张量计算归一化相关性
        else: # 视差级别0
            volume[:, :, i, :, :] = norm_correlation(refimg_fea, targetimg_fea) # 直接计算两个张量的归一化相关性
    volume = volume.contiguous() # 形成归一化相关性体积，输出(B,1,D/4,H/4,W/4)
    return volume

def SpatialTransformer_grid(x, y, disp_range_samples):
    bs, channels, height, width = y.size()
    ndisp = disp_range_samples.size()[1]

    mh, mw = torch.meshgrid([
        torch.arange(0, height, dtype=x.dtype, device=x.device),
        torch.arange(0, width, dtype=x.dtype, device=x.device)]) # (H *W)
    mh = mh.reshape(1, 1, height, width).repeat(bs, ndisp, 1, 1)
    mw = mw.reshape(1, 1, height, width).repeat(bs, ndisp, 1, 1) # (B, D, H, W)

    cur_disp_coords_y = mh
    cur_disp_coords_x = mw - disp_range_samples
    coords_x = cur_disp_coords_x / ((width - 1.0) / 2.0) - 1.0 # trans to -1 - 1
    coords_y = cur_disp_coords_y / ((height - 1.0) / 2.0) - 1.0
    grid = torch.stack([coords_x, coords_y], dim=4) # (B, D, H, W, 2)

    y_warped = F.grid_sample(y, grid.view(bs, ndisp * height, width, 2), mode='bilinear', padding_mode='zeros', align_corners=True).view(bs, channels, ndisp, height, width) # (B, C, D, H, W)

    return y_warped

# 上下文上采样
def context_upsample(depth_low, up_weights): # 输入1/4原尺寸回归视差(1通道)，原尺寸"超像素"权重(9通道)
    ###
    # cv (b,1,h,w)
    # sp (b,9,4*h,4*w)
    ###
    b, c, h, w = depth_low.shape # 1/4原尺寸回归视差(1通道)

    depth_unfold = F.unfold(depth_low.reshape(b, c, h, w), 3, 1, 1).reshape(b, -1, h, w) # 以窗口3*3步长1展开低分辨率视差(得到若干列向量)，改变形状为(B,(3*3),H/4,W/4)
    depth_unfold = F.interpolate(depth_unfold, (h*4, w*4), mode='nearest').reshape(b, 9, h*4, w*4) # 最近邻插值为原尺寸，输出(B,9,H,W)(1/4尺寸视差)

    depth = (depth_unfold * up_weights).sum(1) # 插值视差*上下文权重，在通道维度累和，输出(B,H,W)(1/4尺寸视差)

    return depth

# 视差回归，在视差维度取两个相关性最大值做回归
def regression_topk(cost, disparity_samples, k): # 输入聚合的相关性成本体积(B,D/4,H/4,W/4)，视差序列体积(B,D/4,H/4,W/4)，2
    _, ind = cost.sort(1, True) # 在视差维度对相关性成本体积降序排序，输出排序的索引张量，(B,D/4,H/4,W/4)
    pool_ind = ind[:, :k] # 获得前两个值的索引(值越大左右特征越相关)，(B,2,H/4,W/4)
    cost = torch.gather(cost, 1, pool_ind) # 在视差维度按照指定索引查找相关性成本体积，(B,2,H/4,W/4)
    prob = F.softmax(cost, 1) # 概率体积，在视差维度对前两个相关性值做softmax，(B,2,H/4,W/4)
    disparity_samples = torch.gather(disparity_samples, 1, pool_ind) # 对应视差级别，在视差维度按照指定索引查找视差序列，(B,2,H/4,W/4)
    pred = torch.sum(disparity_samples * prob, dim=1, keepdim=True) # 回归视差，视差级别*概率体积，在视差维度累和，输出(B,1,H/4,W/4)
    return pred
